package mysql

import (
	"fmt"
	"github.com/go-xorm/xorm"
	cmap "github.com/orcaman/concurrent-map"
	"github.com/pkg/errors"
	"regexp"
	"strings"
)

var (
	dbTableList         cmap.ConcurrentMap // 数据库表字段映射 map
	sqlBlackListPattern string             // SQL 黑名单正则表达式
)

func init() {
	dbTableList = cmap.New()
	sqlBlackListPattern = getSqlBlacklistPattern()
}

// Field && Cols 字符串条件白名单防 SQL 注入
func filterColumns(columns []string) error {
	for _, column := range columns {
		if b, err := regexp.MatchString(sqlBlackListPattern, strings.ToLower(column)); err != nil {
			return err
		} else if b {
			return errors.New("abnormal column：" + column)
		}
	}
	return nil
}

// 过滤 order by 条件潜在 SQL 注入条件
// 思路：获取操作表所有字段，替换掉“,”、concat([\s\S]*)、tableName、columnName、“`”、空格，如果剩余字符串还有内容，则说明排序值有问题
func (h *Handler) filterOrderBy(table map[string]string, cond string, fields []string) error {
	// 0. 获取引擎
	var engine *xorm.Engine
	if h.EngineGroup != nil {
		engine = h.EngineGroup.Slave()
	} else {
		engine = h.Engine
	}
	// 1. 获取当前表数据字段列
	list, err := getDatabaseTableList(engine)
	if err != nil {
		return err
	}
	// 2. 声明验证字段集
	columns := make([]string, 0)
	tables := make([]string, 0)
	for k, v := range table {
		tables = append(tables, k+".", v+".")
		if len(list[v]) == 0 {
			return errors.New("filterOrderBy：current table is abnormal, table：" + v)
		}
		columns = append(columns, list[v]...)
	}
	// 3. 判断字段是否为空，不为空则取出指定查询字段
	if len(fields) != 0 {
		fieldColumn := getColumn(fields)
		columns = append(columns, fieldColumn...)
	}
	// 4. 字符串预处理
	str := strings.ToLower(cond)
	str = strings.ReplaceAll(str, "`", "") // 替换去掉“`”
	str = strings.ReplaceAll(str, "concat(", "")
	str = strings.ReplaceAll(str, ")", "")
	str = strings.ReplaceAll(str, `"`, "")
	// 5. 组装正则替换字符串
	columnPattern := strings.Join(columns, "|")
	tablePattern := strings.Join(tables, "|")
	pattern := fmt.Sprintf(`,|\b(%s)\b|\b([%s]+)\b|\b([ |\n|\t]+desc)\b|\b([ |\n|\t]+asc)\b|[ |\n|\t]+`, columnPattern, tablePattern)
	fmt.Println(pattern)
	re, err := regexp.Compile(pattern)
	if err != nil {
		return errors.New(fmt.Sprintf("[order by] regexp error：%s，cond：%s，pattern：%s", err.Error(), cond, pattern))
	}
	// 6. 正则替换掉所有匹配字符，检查结果是否为空
	result := re.ReplaceAllString(str, "")
	fmt.Println(result)
	if result != "" {
		return errors.New("[order by] condition is abnormal：" + cond + "，result：" + result)
	}
	return nil
}

// 去掉表名称
//func clearTable(table map[string]string, str string) string {
//	for k, v := range table {
//		str = strings.ReplaceAll(str, k+".", "")
//		str = strings.ReplaceAll(str, v+".", "")
//	}
//	return str
//}

// 过滤字段
func getColumn(fields []string) []string {
	columns := make([]string, 0)
	for _, column := range fields {
		column = strings.ReplaceAll(column, "`", "") // 替换去掉“`”
		arr := strings.Split(column, " ")
		if len(arr) == 1 && strings.Contains(column, ".") {
			arr = strings.Split(column, ".")
		}
		column = arr[len(arr)-1]
		if strings.Contains(column, "*") {
			continue
		}
		columns = append(columns, column)
	}
	return columns
}

// 获取当前数据库的表名+字段名称，组装 Map
func getDatabaseTableList(engine *xorm.Engine) (map[string][]string, error) {
	// 0. 获取当前库名称
	dbName := getDatabaseName(engine)
	// 1. 判断是否已存在对应库表信息
	if v, ok := dbTableList.Get(dbName); ok {
		return v.(map[string][]string), nil
	}
	// 2.获取数据表列表信息
	list, err := engine.DBMetas()
	if err != nil {
		return nil, err
	}
	// 3. 获取数据库表对应的字段Map
	tableList := make(map[string][]string, 0)
	for _, info := range list {
		box := make([]string, 0)
		for _, table := range info.Columns() {
			box = append(box, table.Name)
		}
		tableList[info.Name] = box
	}
	dbTableList.Set(dbName, tableList)
	// 4. 加入内存管理
	return tableList, nil
}

// 获取数据库名称
func getDatabaseName(engine *xorm.Engine) string {
	str := engine.DataSourceName()
	reg := regexp.MustCompile(`\/[\d\w_]+\?`)
	db := reg.FindString(str)
	db = strings.Replace(db, "/", "", 1)
	db = strings.Replace(db, "?", "", 1)
	return db
}

// Group By 字符串条件防 SQL 注入
// 思路：获取 filed、表中所有字段 column，替换掉 order by 条件中的所有“,”、concat()、tableName.、column、“`”、空格，如果剩余字符串还有内容，则说明排序值有问题
func (h *Handler) filterGroupBy(table, cond string, fields []string) error {
	// 0. 获取引擎
	var engine *xorm.Engine
	if h.EngineGroup != nil {
		engine = h.EngineGroup.Slave()
	} else {
		engine = h.Engine
	}
	// 1. 获取当前查询数据字段列
	list, err := getDatabaseTableList(engine)
	if err != nil {
		return err
	}
	// 2. 声明验证字段集
	columns := list[table]
	if len(columns) == 0 {
		return errors.New("filterGroupBy：current table is abnormal, table：" + table)
	}
	// 3. 判断字段是否为空
	if len(fields) != 0 {
		fieldColumn := getColumn(fields)
		columns = append(columns, fieldColumn...)
		//for _, column := range fields {
		//	column = strings.ReplaceAll(column, "`", "")       // 替换去掉“`”
		//	column = strings.ReplaceAll(column, table+".", "") // 替换去掉表名
		//	arr := strings.Split(column, " ")
		//	columns = append(columns, arr[len(arr)-1])
		//}
	}
	// 4. 分组条件字符串预处理
	str := strings.ToLower(cond)
	str = strings.ReplaceAll(str, "`", "")       // 替换去掉“`”
	str = strings.ReplaceAll(str, table+".", "") // 替换去掉表名
	str = strings.ReplaceAll(str, "concat(", "")
	str = strings.ReplaceAll(str, ")", "")
	str = strings.ReplaceAll(str, `"`, "")
	// 3. 组装正则替换字符串
	columnPattern := strings.Join(columns, "|")
	pattern := fmt.Sprintf(`,|%s|[ |\n|\t]+`, columnPattern)
	re, err := regexp.Compile(pattern)
	if err != nil {
		return err
	}
	// 4. 正则替换掉所有匹配字符，检查结果是否为空
	if re.ReplaceAllString(str, "") != "" {
		return errors.New("[group by] condition is abnormal：" + cond)
	}
	return nil
}

// 过滤字符串中的非法字符，防 SQL 注入
// 思路：黑名单过滤
// 获取 SQL 字符串黑名单过滤正则表达式
func getSqlBlacklistPattern() string {
	// 过滤字符：| & %0a # ; - + /
	patternOne := `(\|)|(&)|(%)|(#)|(;)|(-)|(\+)|(\/)`
	// 过滤 mysql 关键字
	patternTwo := ""
	keyword := []string{
		"select", "insert", "delete", "drop", "update", "truncate",
		"from", "where", "and", "or", "xor", "like", "between",
		"left", "join", "group by", "limit", "union",
		"database", "information_schema", "mid", "substring", "regexp", "load_file", "into", "outfile",
	}
	length := len(keyword)
	for k, v := range keyword {
		if k != 0 && k+1 < length {
			patternTwo += "|"
		}
		patternTwo += fmt.Sprintf("(%s[ |\n|\t]+)", v)
	}
	patternTwo = fmt.Sprintf(`\b(%s)\b`, patternTwo)
	// 过滤 mysql 方法：get_lock()、sleep()、benchmark()、updatexml()、extractvalue()、hex()
	patternThree := ""
	keyword = []string{
		"get_lock", "sleep", "benchmark", "updatexml", "extractvalue", "hex",
	}
	length = len(keyword)
	for k, v := range keyword {
		if k != 0 && k+1 < length {
			patternThree += "|"
		}
		patternThree += fmt.Sprintf(`%s\(`, v)
	}
	patternThree = fmt.Sprintf(`\b(%s)`, patternThree)
	return fmt.Sprintf("%s|(%s)|(%s)", patternOne, patternTwo, patternThree)
}
